import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
from typing import Union
from torch import Tensor, nn
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer


class HFEmbedder(nn.Module):
    def __init__(self, model_path: str, max_length: int, model_type: str, device: str = "npu"):
        super().__init__()
        # self.is_clip = version.startswith("openai")
        # self.max_length = max_length
        # self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
        # self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(model_path, max_length=max_length)
        # self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(model_path, **hf_kwargs)
        # if self.is_clip:
        # self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
        # self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
        # else:
        # self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length)
        # self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
        # raise ValueError("Version does not start with 'openai', which is required for CLIP models. If you intend to use T5, please adjust the version parameter accordingly.")

        self.model_path = model_path
        self.max_length = max_length
        self.model_type = model_type
        self.device = device

        if model_type == "CLIP":
            self.tokenizer = CLIPTokenizer.from_pretrained(model_path)
            self.hf_module = CLIPTextModel.from_pretrained(model_path)
            self.output_key = "pooler_output"
        elif model_type == "T5":
            self.tokenizer = T5Tokenizer.from_pretrained(model_path)
            self.hf_module = T5EncoderModel.from_pretrained(model_path)
            self.output_key = "last_hidden_state"  # T5模型的输出键
        else:
            raise ValueError(f"Unsupported model type: {model_type}. Choose 'CLIP' or 'T5'.")

        self.hf_module = self.hf_module.eval().requires_grad_(False)
        # self.hf_module = self.hf_module.eval().requires_grad_(False)
        self.hf_module.to(device).to(torch.float32)

    def forward(self, text: list[str]) -> Tensor:
        batch_encoding = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            return_length=False,
            return_overflowing_tokens=False,
            padding="max_length",
            return_tensors="pt",
        )

        outputs = self.hf_module(
            input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
            attention_mask=None,
            output_hidden_states=False,
        )
        embedded_text = outputs[self.output_key]  # 添加代码、获取嵌入表示
        print(f"Embedded text shape: {embedded_text.shape}")  # 打印嵌入的维度
        return embedded_text
        # return outputs[self.output_key]
